'''
数据集预处理脚本cataractDataset
分类数据集
'''
import sys, os
import re
import json
from data_preprocess.utils.crop import crop_resize_save
from data_preprocess.utils.stat import save_classification_stats
from tqdm import tqdm
def infer_category(folder_name):
    """
    根据文件夹名称推断诊断文本。
    规则：
      - 如果名称中包含 "normal"（不区分大小写），则返回 "normal"
      - 如果名称中包含 "cataract"，则返回 "cataract"
      - 如果名称中包含 "glaucoma"，则返回 "glaucoma"
      - 如果名称中包含 "retina"，则返回 "retina disease"
      - 否则抛出异常
    """
    lower = folder_name.lower()
    if "normal" in lower:
        return "normal"
    elif "cataract" in lower:
        return "cataract"
    elif "glaucoma" in lower:
        return "glaucoma"
    elif "retina" in lower:
        return "retina disease"
    else:
        raise ValueError(f"无法从文件夹名称中推断类别：{folder_name}")

def gather_data(data_path, tar_path, prefix='cataract_', resize=(512,512)):
    """
    对新数据集进行预处理：
      - 数据集根目录中包含若干类别文件夹，其路径可能直接位于 data_path 下，
        或位于 data_path/dataset 下。
      - 根据文件夹名称推断每个类别的诊断文本。
      - 遍历每个类别文件夹中的所有图像，调用 crop_resize_save 进行裁剪与 resize，
        并统一保存到 tar_path/images 下，新图像名称格式为
            {prefix}_{原始文件名}.png
        （保留原始文件名，仅修改后缀并加上前缀）。
      - 生成标注信息字典，格式如下：
            {
                new_image_name: {
                    "image_name": new_image_name,
                    "image_path": "images/{new_image_name}",
                    "diagnosis": {"text": <诊断文本>}
                },
                ...
            }
      - 将标注信息保存到 tar_path/annotations.json。
    
    参数：
        data_path (str): 数据集根目录，要求包含类似如下的目录：
                     data_path/1_normal/
                     data_path/2_cataract/
                     data_path/dataset/2_glaucoma/
                     data_path/3_retina_disease/
        tar_path (str): 预处理后数据存放目录。
        prefix (str): 处理后图片名称的前缀，默认为 'eyepacs_'。
        resize (tuple): 目标尺寸，默认为 (512,512)
    
    返回：
        dict: 标注信息字典。
    """
    os.makedirs(tar_path, exist_ok=True)
    images_out_dir = os.path.join(tar_path, "images")
    os.makedirs(images_out_dir, exist_ok=True)
    
    annotations = {}
    
    # 收集所有类别文件夹：包括直接位于 data_path 下的和 data_path/dataset 下的
    category_dirs = []
    for d in os.listdir(data_path):
        full_path = os.path.join(data_path, d)
        if os.path.isdir(full_path) and re.match(r"^\d+_.*", d):
            category_dirs.append(full_path)
    dataset_dir = os.path.join(data_path, "dataset")
    if os.path.exists(dataset_dir) and os.path.isdir(dataset_dir):
        for d in os.listdir(dataset_dir):
            full_path = os.path.join(dataset_dir, d)
            if os.path.isdir(full_path) and re.match(r"^\d+_.*", d):
                category_dirs.append(full_path)
    
    if not category_dirs:
        raise ValueError("未找到符合格式的类别文件夹，请检查数据集目录结构。")
    
    # 遍历每个类别文件夹
    for cat_dir in category_dirs:
        folder_name = os.path.basename(cat_dir)
        try:
            diag_text = infer_category(folder_name)
        except ValueError as e:
            raise ValueError(f"在处理文件夹 '{cat_dir}' 时出错：{e}")
        
        for image_file in sorted(os.listdir(cat_dir)):
            if not image_file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.ppm')):
                continue
            src_image_path = os.path.join(cat_dir, image_file)
            base_name = os.path.splitext(image_file)[0]
            new_image_name = f"{prefix}_{base_name}.png"
            dest_image_path = os.path.join(images_out_dir, new_image_name)
            
            crop_info=crop_resize_save(
                image_path=src_image_path,
                save_path=dest_image_path,
                resize=resize,
                crop_threshold=25
            )
            
            annotations[new_image_name] = {
                "image_name": new_image_name,
                "image_path": dest_image_path,
                'original_path':src_image_path,
                'crop_info':crop_info,
                "diagnosis": {
                    "classification": {
                        "text": diag_text
                    }
                }
            }
    # 统计分类数据分布
    save_classification_stats(annotations, './experiments/stat/catact.json')
    # 保存标注信息
    annotations_path = os.path.join(tar_path, "annotations.json")
    with open(annotations_path, "w", encoding="utf-8") as f:
        json.dump(annotations, f, indent=4)
    
    return annotations

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="新数据集预处理，根据目录名称推断类别")
    parser.add_argument("--data_path", type=str, default="/home/zhangpinglu/data0/gy/Dataset/public_dataset/cataractDataset",
                        help="数据集根目录，要求包含如：1_normal, 2_cataract, dataset/2_glaucoma, 3_retina_disease 等目录")
    parser.add_argument("--tar_path", type=str, default="/home/zhangpinglu/data0/gy/Dataset/public_processed/cataract",
                        help="预处理后数据存放目录")
    parser.add_argument("--prefix", type=str, default="cataract",
                        help="处理后图片名称的前缀，默认为 cataract")
    parser.add_argument("--resize", type=int, nargs=2, default=[512,512],
                        help="目标尺寸，默认为 512 512")
    args = parser.parse_args()
    print("本数据仅有601张数据，因此不设置进度条（tqdm），请稍后...")
    annotations = gather_data(args.data_path, args.tar_path, prefix=args.prefix, resize=tuple(args.resize))
    print("Preprocessing completed. Annotations saved.")
